import os
import requests
from PIL import Image
import numpy as np
import torch
from io import BytesIO

class XP_Utils_Image_Concat:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {"required":
                {
                    "image_urls": ("STRING", {"default": ""}),
                    "delimiter": ("STRING", {"default": ","}),
                    "max_width": ("INT", {"default": 2048, "min": 64, "max": 8192, "step": 64}),
                },
                "optional": {
                    "spacing": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}),
                }
               }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "sample"
    CATEGORY = "XP_Utils"

    def sample(self, image_urls, delimiter=",", max_width=2048, spacing=0):
        # 分割图片URL
        urls = [url.strip() for url in image_urls.split(delimiter) if url.strip()]
        
        if not urls:
            # 如果没有有效的URL，返回一个空白图片
            empty_image = np.zeros((64, 64, 3), dtype=np.float32)
            return (torch.from_numpy(empty_image.astype(np.float32) / 255.0).unsqueeze(0),)
        
        # 加载所有图片
        images = []
        for url in urls:
            try:
                # 尝试从本地文件加载
                if os.path.exists(url):
                    img = Image.open(url).convert("RGB")
                # 尝试从URL加载
                else:
                    response = requests.get(url)
                    img = Image.open(BytesIO(response.content)).convert("RGB")
                images.append(img)
            except Exception as e:
                print(f"无法加载图片 {url}: {e}")
        
        if not images:
            # 如果没有成功加载任何图片，返回一个空白图片
            empty_image = np.zeros((64, 64, 3), dtype=np.float32)
            return (torch.from_numpy(empty_image.astype(np.float32) / 255.0).unsqueeze(0),)
        
        # 计算总宽度和最大高度
        total_width = sum(img.width for img in images) + spacing * (len(images) - 1)
        max_height = max(img.height for img in images)
        
        # 如果总宽度超过最大宽度，按比例缩小所有图片
        scale = 1.0
        if total_width > max_width:
            scale = max_width / total_width
            scaled_images = []
            for img in images:
                new_width = int(img.width * scale)
                new_height = int(img.height * scale)
                scaled_img = img.resize((new_width, new_height), Image.LANCZOS)
                scaled_images.append(scaled_img)
            images = scaled_images
            total_width = sum(img.width for img in images) + spacing * (len(images) - 1)
            max_height = max(img.height for img in images)
        
        # 创建一个新的空白图片
        result = Image.new('RGB', (total_width, max_height))
        
        # 将所有图片横向拼接
        x_offset = 0
        for img in images:
            result.paste(img, (x_offset, (max_height - img.height) // 2))
            x_offset += img.width + spacing
        
        # 将PIL图像转换为ComfyUI可用的格式（torch tensor）
        result_np = np.array(result).astype(np.float32) / 255.0
        result_tensor = torch.from_numpy(result_np).unsqueeze(0)
        
        return (result_tensor,)
